# -*- coding: UTF-8 -*-

from pathlib import Path, WindowsPath, PosixPath
import numpy as np

import json


def create_txt(text, name, save_path="../generate_dir"):
    txt_dir = "txt"
    save_path = Path(save_path) / txt_dir
    if not save_path.exists():
        save_path.mkdir(parents=True)

    save_path = save_path / "txt.jsonl"

    fp = save_path.open("a+", encoding="utf-8")
    if len(text) == 0:
        text = "这是一张背景图,无文字"

    else:
        text = "这是一张带有文字的图片,图片中的文字如下:" + text
    save_text = json.dumps({"img_file": str(name.absolute()), "content": text}) + "\n"
    fp.write(save_text)
    fp.close()


def _get_all_file(_base_path) -> list:
    """
    获取文件夹下的所有文件
    :param _base_path: 基础路径, 如果是文件则会返回文件的list
    :return: r
    """
    _base_path = Path(_base_path)

    if _base_path.is_file():
        return [_base_path]

    _dirs = [_base_path]
    _files = []
    while True:
        if len(_dirs) == 0:
            break

        _dir = _dirs.pop(0)

        for _file in _dir.iterdir():
            if _file.is_file():
                _files.append(_file)
            else:
                _dirs.append(_file)

    return _files


def _filter_background_img(_base_path, _num):
    """
    筛选背景图片
    :param _base_path: 图片父目录
    :param _num: 筛选数量
    :return: r
    """
    _files = _get_all_file(_base_path)
    if len(_files) < _num:
        raise ValueError("selectNum:{} greater than filesNum:{}".format(_num, len(_files)))
    _filter_files = np.random.choice(_files, replace=False, size=_num)

    return _filter_files


def select_background_img(_base_path, _num) -> list[WindowsPath or PosixPath]:
    """
    筛选背景图片
    :param _base_path: 背景图片父目录
    :param _num: 筛选数量
    :return: 筛选路径集合
    """
    path = Path(_base_path)
    if not path.exists():
        raise FileExistsError("dir:[{}] not exist!".format(path.absolute()))

    return _filter_background_img(_base_path, _num)


def generate_background_image(img_folder_path, save_path="../generate_dir/uu", _num=20000):
    back_files = select_background_img(img_folder_path, _num)
    for _file in back_files:
        text = ""
        create_txt(text, _file, save_path)
